import os
import torch.nn as nn

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import numpy as np, argparse, time
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler
from meld_data_loader import MELDDataset
from meld_model import Hierarchical_Transformer_Model
from loss_function  import MaskedNLLLoss, MaskedKLDivLoss
from sklearn.metrics import f1_score, confusion_matrix, accuracy_score, classification_report
import pickle as pk
import datetime


# 定义获取训练和验证集采样器的函数，参数为训练集，验证比例，数据集
def get_train_valid_sampler(trainset, valid=0.1, dataset='MELD'):
    # 获取训练集大小
    size = len(trainset)
    # 创建一个从0到size的列表
    idx = list(range(size))
    # 计算验证集大小
    split = int(valid * size)
    # 返回从split到size的随机采样器，从0到split的随机采样器
    return SubsetRandomSampler(idx[split:]), SubsetRandomSampler(idx[:split])


# 定义获取MELD数据加载器的函数，参数为batch_size，验证比例，num_workers，pin_memory
def get_MELD_loaders(batch_size=32, valid=0.1, num_workers=0, pin_memory=False):
    # 创建MELD数据集
    trainset = MELDDataset('data/meld_multimodal_features.pkl')
    # 获取训练集和验证集采样器
    train_sampler, valid_sampler = get_train_valid_sampler(trainset, valid, 'MELD')
    # 创建训练集加载器
    train_loader = DataLoader(trainset,
                              batch_size=batch_size,
                              sampler=train_sampler,
                              collate_fn=trainset.collate_fn,
                              num_workers=num_workers,
                              pin_memory=pin_memory)
    # 创建验证集加载器
    valid_loader = DataLoader(trainset,
                              batch_size=batch_size,
                              sampler=valid_sampler,
                              collate_fn=trainset.collate_fn,
                              num_workers=num_workers,
                              pin_memory=pin_memory)

    # 创建测试集
    testset = MELDDataset('data/meld_multimodal_features.pkl', train=False)
    # 创建测试集加载器
    test_loader = DataLoader(testset,
                             batch_size=batch_size,
                             collate_fn=testset.collate_fn,
                             num_workers=num_workers,
                             pin_memory=pin_memory)
    # 返回训练集加载器，验证集加载器，测试集加载器
    return train_loader, valid_loader, test_loader


# 定义一个函数，用于获取IEMOCAP加载器，参数包括batch_size（默认值为32），valid（默认值为0.1），num_workers（默认值为0），pin_memory（默认值为False）
def get_IEMOCAP_loaders(batch_size=32, valid=0.1, num_workers=0, pin_memory=False):
    # 创建IEMOCAP数据集
    trainset = IEMOCAPDataset()
    # 获取训练集和验证集的采样器
    train_sampler, valid_sampler = get_train_valid_sampler(trainset, valid)
    # 创建训练集加载器
    train_loader = DataLoader(trainset,
                              batch_size=batch_size,
                              sampler=train_sampler,
                              collate_fn=trainset.collate_fn,
                              num_workers=num_workers,
                              pin_memory=pin_memory)
    # 创建验证集加载器
    valid_loader = DataLoader(trainset,
                              batch_size=batch_size,
                              sampler=valid_sampler,
                              collate_fn=trainset.collate_fn,
                              num_workers=num_workers,
                              pin_memory=pin_memory)

    # 创建测试集
    testset = IEMOCAPDataset(train=False)
    # 创建测试集加载器
    test_loader = DataLoader(testset,
                             batch_size=batch_size,
                             collate_fn=testset.collate_fn,
                             num_workers=num_workers,
                             pin_memory=pin_memory)
    # 返回训练集加载器，验证集加载器，测试集加载器
    return train_loader, valid_loader, test_loader


def train_or_eval_model(model, loss_function, kl_loss, dataloader, epoch, optimizer=None, train=False, gamma_1=1.0,
                        gamma_2=1.0, gamma_3=1.0, mode=3):
    losses, preds, labels, masks = [], [], [], []

    assert not train or optimizer != None
    if train:
        model.train()
    else:
        model.eval()

    for data in dataloader:
        if train:
            optimizer.zero_grad()

        # 将data中的数据转换为cuda格式  qmask是说话人
        textf, visuf, acouf, qmask, umask, label = [d.cuda() for d in data[:-1]] if cuda else data[:-1]
        # 将qmask的形状转换为(batch_size, n_heads, query_len, key_len)
        qmask = qmask.permute(1, 0, 2)
        # 计算每个元素的长度
        lengths = [(umask[j] == 1).nonzero().tolist()[-1][0] + 1 for j in range(len(umask))]

        log_prob1, log_prob2, log_prob3, all_log_prob, all_prob, \
        kl_log_prob1, kl_log_prob2, kl_log_prob3, kl_all_prob = model(textf, visuf, acouf, umask, qmask, lengths, mode)

        lp_1 = log_prob1.view(-1, log_prob1.size()[2])
        lp_2 = log_prob2.view(-1, log_prob2.size()[2])
        lp_3 = log_prob3.view(-1, log_prob3.size()[2])
        lp_all = all_log_prob.view(-1, all_log_prob.size()[2])
        labels_ = label.view(-1)

        kl_lp_1 = kl_log_prob1.view(-1, kl_log_prob1.size()[2])
        kl_lp_2 = kl_log_prob2.view(-1, kl_log_prob2.size()[2])
        kl_lp_3 = kl_log_prob3.view(-1, kl_log_prob3.size()[2])
        kl_p_all = kl_all_prob.view(-1, kl_all_prob.size()[2])

        # NLL损失函数  70
        # nll_loss = nn.NLLLoss()
        # loss = nll_loss(lp_all, labels_)

        # 交叉熵损失函数  69
        criterion = nn.CrossEntropyLoss()
        loss = criterion(lp_all, labels_)

        # 修改去掉蒸馏模型
        # loss = gamma_1 * loss_function(lp_all, labels_, umask)

        # 计算损失函数
        # loss = gamma_1 * loss_function(lp_all, labels_, umask) + \
        #         gamma_2 * (loss_function(lp_1, labels_, umask) + loss_function(lp_2, labels_, umask) + loss_function(lp_3, labels_, umask)) + \
        #        gamma_3 * (kl_loss(kl_lp_1, kl_p_all, umask) + kl_loss(kl_lp_2, kl_p_all, umask) + kl_loss(kl_lp_3, kl_p_all, umask))

        lp_ = all_prob.view(-1, all_prob.size()[2])

        # 获取每个像素的类别
        pred_ = torch.argmax(lp_, 1)
        # 将每个像素的类别添加到preds列表中
        preds.append(pred_.data.cpu().numpy())
        # 将真实类别添加到labels列表中
        labels.append(labels_.data.cpu().numpy())
        # 将每个像素的mask添加到masks列表中
        masks.append(umask.view(-1).cpu().numpy())
        losses.append(loss.item() * masks[-1].sum())
        if train:
            loss.backward()
            if args.tensorboard:
                for param in model.named_parameters():
                    if param.grad is not None:
                        writer.add_histogram(param[0], param[1].grad, epoch)
            optimizer.step()

    # 计算模型在验证集上的指标
    if preds != []:
        preds = np.concatenate(preds)
        labels = np.concatenate(labels)
        masks = np.concatenate(masks)
    else:
        return float('nan'), float('nan'), [], [], [], float('nan')

    # 计算模型的损失、准确率和f1分数
    avg_loss = round(np.sum(losses) / np.sum(masks), 4)
    avg_accuracy = round(accuracy_score(labels, preds, sample_weight=masks) * 100, 2)
    avg_fscore = round(f1_score(labels, preds, sample_weight=masks, average='weighted') * 100, 2)
    return avg_loss, avg_accuracy, labels, preds, masks, avg_fscore


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--no-cuda', action='store_true', default=False, help='does not use GPU')
    parser.add_argument('--lr', type=float, default=0.0001, metavar='LR', help='learning rate')
    parser.add_argument('--l2', type=float, default=0.00001, metavar='L2', help='L2 regularization weight')
    parser.add_argument('--dropout', type=float, default=0.5, metavar='dropout', help='dropout rate')
    parser.add_argument('--batch-size', type=int, default=16, metavar='BS', help='batch size')
    parser.add_argument('--hidden_dim', type=int, default=1024, metavar='hidden_dim', help='output hidden size')
    parser.add_argument('--n_head', type=int, default=8, metavar='n_head', help='number of heads')
    parser.add_argument('--epochs', type=int, default=150, metavar='E', help='number of epochs')
    parser.add_argument('--temp', type=int, default=1, metavar='temp', help='temp')
    parser.add_argument('--tensorboard', action='store_true', default=False, help='Enables tensorboard log')
    parser.add_argument('--class-weight', action='store_true', default=True, help='use class weights')
    parser.add_argument('--Dataset', default='IEMOCAP', help='dataset to train and test')
    parser.add_argument('--gamma_1', type=float, default=1.0, help='proportion')
    parser.add_argument('--gamma_2', type=float, default=1.0, help='proportion')
    parser.add_argument('--gamma_3', type=float, default=1.0, help='proportion')
    parser.add_argument('--num_block', type=int, default=2, help='Working mode')
    parser.add_argument('--mode', type=str, default='so', help='mode')

    args = parser.parse_args()
    today = datetime.datetime.now()
    print(args)

    args.cuda = torch.cuda.is_available() and not args.no_cuda
    if args.cuda:
        print('Running on GPU')
    else:
        print('Running on CPU')

    if args.tensorboard:
        from tensorboardX import SummaryWriter

        writer = SummaryWriter()

    cuda = args.cuda
    n_epochs = args.epochs
    batch_size = args.batch_size
    feat2dim = {'IS10': 1582, 'denseface': 342, 'MELD_audio': 300}
    D_audio = feat2dim['IS10'] if args.Dataset == 'IEMOCAP' else feat2dim['MELD_audio']
    D_visual = feat2dim['denseface']
    D_text = 1024

    D_m = D_audio + D_visual + D_text

    n_speakers = 9 if args.Dataset == 'MELD' else 2
    n_classes = 7 if args.Dataset == 'MELD' else 6 if args.Dataset == 'IEMOCAP' else 1

    print('temp {}'.format(args.temp))

    model = Hierarchical_Transformer_Model(args.Dataset, args.temp, D_text, D_visual, D_audio, args.n_head,
                                    n_classes=n_classes,
                                    hidden_dim=args.hidden_dim,
                                    n_speakers=n_speakers, num_block=args.num_block,
                                    dropout=args.dropout)

    total_params = sum(p.numel() for p in model.parameters())
    print('total parameters: {}'.format(total_params))
    total_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('training parameters: {}'.format(total_trainable_params))

    if cuda:
        model.cuda()

    kl_loss = MaskedKLDivLoss()
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2)

    if args.Dataset == 'MELD':
        loss_function = MaskedNLLLoss()
        train_loader, valid_loader, test_loader = get_MELD_loaders(valid=0.0,
                                                                   batch_size=batch_size,
                                                                   num_workers=0)
    elif args.Dataset == 'IEMOCAP':
        # 获取IEMOCAP数据集的损失权重
        loss_weights = torch.FloatTensor([1 / 0.086747,
                                          1 / 0.144406,
                                          1 / 0.227883,
                                          1 / 0.160585,
                                          1 / 0.127711,
                                          1 / 0.252668])
        # 获取损失函数，并使用损失权重
        loss_function = MaskedNLLLoss(loss_weights.cuda() if cuda else loss_weights)

        # loss_function = MaskedNLLLoss()

        # 获取IEMOCAP数据集的加载器
        train_loader, valid_loader, test_loader = get_IEMOCAP_loaders(valid=0.0,
                                                                      batch_size=batch_size,
                                                                      num_workers=0)
    else:
        print("There is no such dataset")

    best_fscore, best_loss, best_label, best_pred, best_mask = None, None, None, None, None
    all_fscore, all_acc, all_loss, tr_loss, te_loss = [], [], [], [],[]

    for e in range(n_epochs):
        start_time = time.time()

        train_loss, train_acc, _, _, _, train_fscore = train_or_eval_model(model, loss_function, kl_loss, train_loader,
                                                                           e, optimizer, True, gamma_1=args.gamma_1,
                                                                           gamma_2=args.gamma_2, gamma_3=args.gamma_3,
                                                                           mode=args.mode)
        valid_loss, valid_acc, _, _, _, valid_fscore = train_or_eval_model(model, loss_function, kl_loss, valid_loader,
                                                                           e, gamma_1=args.gamma_1,
                                                                           gamma_2=args.gamma_2, gamma_3=args.gamma_3,
                                                                           mode=args.mode)
        test_loss, test_acc, test_label, test_pred, test_mask, test_fscore = train_or_eval_model(model, loss_function,
                                                                                                 kl_loss, test_loader,
                                                                                                 e,
                                                                                                 gamma_1=args.gamma_1,
                                                                                                 gamma_2=args.gamma_2,
                                                                                                 gamma_3=args.gamma_3,
                                                                                                 mode=args.mode)
        all_fscore.append(test_fscore)

        if best_fscore == None or best_fscore < test_fscore:
            best_fscore = test_fscore
            best_label, best_pred, best_mask = test_label, test_pred, test_mask

        if args.tensorboard:
            writer.add_scalar('test: accuracy', test_acc, e)
            writer.add_scalar('test: fscore', test_fscore, e)
            writer.add_scalar('train: accuracy', train_acc, e)
            writer.add_scalar('train: fscore', train_fscore, e)

        print(
            'epoch: {}, train_loss: {}, train_acc: {}, train_fscore: {}, valid_loss: {}, valid_acc: {}, valid_fscore: {}, test_loss: {}, test_acc: {}, test_fscore: {}, time: {} sec'. \
            format(e + 1, train_loss, train_acc, train_fscore, valid_loss, valid_acc, valid_fscore, test_loss, test_acc,
                   test_fscore, round(time.time() - start_time, 2)))
        if (e + 1) % 10 == 0:
            print(classification_report(best_label, best_pred, sample_weight=best_mask, digits=4))
            print(confusion_matrix(best_label, best_pred, sample_weight=best_mask))

    if args.tensorboard:
        writer.close()

    print('Test performance..')
    print('F-Score: {}'.format(max(all_fscore)))
    print('F-Score-index: {}'.format(all_fscore.index(max(all_fscore)) + 1))

    if not os.path.exists("IEMOCAP_Exp_Results_{}_{}_{}.pk".format(today.year, today.month, today.day)):
        with open("record_{}_{}_{}.pk".format(today.year, today.month, today.day), 'wb') as f:
            pk.dump({}, f)
    with open("IEMOCAP_Exp_Results_{}_{}_{}.pk".format(today.year, today.month, today.day), 'rb') as f:
        record = pk.load(f)
    key_ = 'name_'
    if record.get(key_, False):
        record[key_].append(max(all_fscore))
    else:
        record[key_] = [max(all_fscore)]
    if record.get(key_ + 'Exp_Re', False):
        record[key_ + 'Exp_Re'].append(classification_report(best_label, best_pred, sample_weight=best_mask, digits=4))
    else:
        record[key_ + 'Exp_Re'] = [classification_report(best_label, best_pred, sample_weight=best_mask, digits=4)]
    with open("Exp_Re_{}_{}_{}.pk".format(today.year, today.month, today.day), 'wb') as f:
        pk.dump(record, f)

    print(classification_report(best_label, best_pred, sample_weight=best_mask, digits=4))
    print(confusion_matrix(best_label, best_pred, sample_weight=best_mask))


